# #
# # 一个服务化的NiiSegmentor类,用于在Python代码中直接调用TotalSegmentator模型,
# # 实现对CT影像的全自动、多器官分割。

# import os
# import nibabel as nib
# import numpy as np

# from totalsegmentator.python_api import totalsegmentator

# class NiiSegmentor:
#     """
#     一个封装了TotalSegmentator功能的分割服务类。
#     """
#     def __init__(self, input_nii_path: str):
#         """
#         初始化NiiSegmentor实例。

#         :param input_nii_path: 输入的Nii文件, 待分割的原始CT影像 (.nii或.nii.gz) 的完整路径
#         """
#         if not os.path.exists(input_nii_path):
#             raise FileNotFoundError(f"输入文件未找到: {input_nii_path}")

#         self.input_nii_path = input_nii_path
#         self.segmentation_result = None

#         print(f"初始化NiiSegmentor, 输入文件: {self.input_nii_path}")

#     def segment_all(self, output_dir: str, fast: bool = False, device: str = "cuda") -> dict:
#         """
#         执行全自动分割。

#         :param output_dir: 输出目录, 分割结果将保存到此目录下
#         :param fast: 是否使用快速模式, 默认为False
#         :param device: 设备类型, 默认为"gpu", 可选"cpu"或"gpu"
#         :return: 包含每个器官名称和对应掩码文件路径的字典
#         """
#         print("\n--- 开始执行分割 ---")
#         print(f"输出目录: {output_dir}")
#         print(f"快速模式: {fast}")
#         print(f"使用设备: {device}")

#         os.makedirs(output_dir, exist_ok=True)

#         try:
#             # [修正点 3] 使用官方API参数名 `output`
#             totalsegmentator(
#                 input=self.input_nii_path,
#                 output=output_dir,
#                 fast=fast,
#                 device=device
#             )
#             print("--- TotalSegmentator 分割任务执行成功 ---")

#         except Exception as e:
#             print(f"分割任务执行失败: {e}")
#             raise RuntimeError(f"分割任务执行失败: {e}")

#         print("--- 整理分割结果 ---")
#         result_dict = {}
#         for file_name in os.listdir(output_dir):
#             if file_name.endswith(('.nii', '.nii.gz')):
#                 base_name = os.path.basename(file_name)
#                 organ_name, _ = os.path.splitext(base_name)
#                 if organ_name.endswith('.nii'):
#                     organ_name, _ = os.path.splitext(organ_name)
                
#                 file_path = os.path.join(output_dir, file_name)
#                 result_dict[organ_name] = file_path

#         self.segmentation_result = result_dict

#         if not self.segmentation_result:
#             raise RuntimeError("分割结果为空，请检查TotalSegmentator的输出和输入文件。")
#         else:
#             print(f"成功分割并整理了 {len(self.segmentation_result)} 个结构。")

#         return self.segmentation_result

#     def get_available_organs(self) -> list:
#         """
#         获取可分割的器官列表。

#         :return: 可分割器官名称列表
#         """
#         if self.segmentation_result is None:
#             raise RuntimeError("请先成功执行 .segment_all() 方法以获取分割结果。")
        
#         return list(self.segmentation_result.keys())

#     def load_mask(self, organ_name: str) -> np.ndarray:
#         """
#         加载指定器官的分割掩码。

#         :param organ_name: 器官名称
#         :return: 分割掩码的numpy数组
#         """
#         if self.segmentation_result is None:
#             raise RuntimeError("请先成功执行 .segment_all() 方法以获取分割结果。")
        
#         if organ_name not in self.segmentation_result:
#             raise ValueError(f"器官 '{organ_name}' 不在分割结果中。可用的器官有: {self.get_available_organs()}")

#         mask_path = self.segmentation_result[organ_name]
#         if not os.path.exists(mask_path):
#             raise FileNotFoundError(f"掩码文件在整理后丢失: {mask_path}")
        
#         mask_nii = nib.load(mask_path)
#         mask_data = mask_nii.get_fdata().astype(np.uint8)
        
#         return mask_data
    


"""
分割重建完整步骤的第二步，将nii.gz完整的文件使用totalsegmentator进行分割，得到每个器官单独得nii.gz文件
"""

import os
import nibabel as nib
import numpy as np
from totalsegmentator.python_api import totalsegmentator
from typing import Optional

class NiiSegmentor:
    def __init__(self, input_nii_path: str):
        if not os.path.exists(input_nii_path):
            raise FileNotFoundError(f"输入文件未找到: {input_nii_path}")
        self.input_nii_path = input_nii_path
        self.segmentation_result=None
        print(f"初始化NiiSegmentor, 输入文件: {self.input_nii_path}")


    def segment(self, output_dir: str, task:str="total",fast: bool = False, device: str = "gpu") -> dict:
        """
        执行全自动分割。
        :param output_dir: 输出目录, 分割结果将保存到此目录下
        :param task: 分割任务类型, 默认为"total", 可选"total"或指定分割任务
        :param fast: 是否使用快速模式, 默认为False
        :param device: 设备类型
        :return: 包含每个器官名称和对应掩码文件路径的字典
        """

        print("\n--- 开始执行分割 ---")
        print(f"输出目录: {output_dir}")
        print(f"分割任务: {task}")
        print(f"快速模式: {fast}")
        print(f"使用设备: {device}")

        os.makedirs(output_dir, exist_ok=True)

        try:
            totalsegmentator(
                input=self.input_nii_path,
                output=output_dir,
                task=task,
                fast=fast,
                device=device
            )
            print("  分割任务执行成功 ")

        except Exception as e:
            print(f"分割任务执行失败: {e}")
            raise RuntimeError(f"分割任务执行失败: {e}")

        print("--- 整理分割结果 ---")
        result_dict = {}
        for file_name in os.listdir(output_dir):
            if file_name.endswith(('.nii', '.nii.gz')):
                organ_name = file_name.split('.nii')[0]
                file_path = os.path.join(output_dir, file_name)
                result_dict[organ_name] = file_path

        self.segmentation_result = result_dict

        if not self.segmentation_result:
            raise RuntimeError("分割结果为空，请检查TotalSegmentator的输出和输入文件。")
        else:
            print(f"成功分割并整理了 {len(self.segmentation_result)} 个结构。")

        return self.segmentation_result
    

    def get_available_organs(self) -> list:
        """
        获取可分割的器官列表。

        :return: 可分割器官名称列表
        """
        if self.segmentation_result is None:
            raise RuntimeError("请先成功执行 .segment_all() 方法以获取分割结果。")
        
        return list(self.segmentation_result.keys())

    def load_mask(self, organ_name: str) -> np.ndarray:
        """
        加载指定器官的分割掩码。

        :param organ_name: 器官名称
        :return: 分割掩码的numpy数组
        """
        if self.segmentation_result is None:
            raise RuntimeError("请先成功执行 .segment_all() 方法以获取分割结果。")
        
        if organ_name not in self.segmentation_result:
            raise ValueError(f"器官 '{organ_name}' 不在分割结果中。可用的器官有: {self.get_available_organs()}")

        mask_path = self.segmentation_result[organ_name]
        if not os.path.exists(mask_path):
            raise FileNotFoundError(f"掩码文件在整理后丢失: {mask_path}")
        
        mask_nii = nib.load(mask_path)  #加载头信息、创建指向文件指针
        mask_data = mask_nii.get_fdata().astype(np.uint8)   #加载数据，把浮点数转变为uint8，从八个字节到一个字节。Mask，值由0，1构成
        
        return mask_data
    




if __name__ == '__main__':
    BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    #指定输入（nii）路径和分割结果的路径，输入是一个完整的nii.gz有效文件，输出是分割后器官的nii.gz文件
    NIFTI_INPUT_DIR = os.path.join(BASE_DIR, "data\\nifti_output\patient001\Abdomen-V  2.0  B30f")
    SEG_OUTPUT_DIR = os.path.join(BASE_DIR, "data\\segmentation_results\\patient001\Abdomen-V  2.0  B30f")
    try:
        print(f"--- 正在智能查找主NIfTI文件于: {NIFTI_INPUT_DIR} ---")
        nifti_files = [f for f in os.listdir(NIFTI_INPUT_DIR) if f.endswith('.nii.gz')]
        
        if not nifti_files:
            raise FileNotFoundError("在输入目录中未找到任何 .nii.gz 文件。")
        
        # 筛选掉带有常见后缀的次要文件
        primary_files = [f for f in nifti_files if '_Eq_' not in f and '_ph' not in f and '_Scout' not in f]
        
        if len(primary_files) == 1:
            main_nifti_filename = primary_files[0]
            print(f"成功找到主文件: {main_nifti_filename}")
        elif len(primary_files) > 1:
            print(f"警告: 找到多个可能的主文件: {primary_files}。将使用第一个。")
            main_nifti_filename = primary_files[0]
        elif len(nifti_files) >= 1: # 如果筛选后没有了，就退回到使用第一个找到的
            print(f"警告: 未能明确区分主文件。将使用找到的第一个文件: {nifti_files[0]}")
            main_nifti_filename = nifti_files[0]
        else:
             raise FileNotFoundError("处理后未找到任何可用的 .nii.gz 文件。")
            
        NIFTI_INPUT_PATH = os.path.join(NIFTI_INPUT_DIR, main_nifti_filename)

        # 执行分割
        segmentor = NiiSegmentor(NIFTI_INPUT_PATH)
        # fast=False提高分割质量
        segmentation_files = segmentor.segment(SEG_OUTPUT_DIR, fast=False)
        print("\n--- 分割完成 ---")
        print("部分结果:")
        for organ, path in list(segmentation_files.items())[:5]:
            print(f"  - {organ}: {path}")
            
    except Exception as e:
        print(f"主程序出错: {e}")